#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
import cv2 as cv
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from ai_library.components.utils import putText

# 设定绿色阈值
lower_green = np.array([38, 43, 43])
upper_green = np.array([80, 255, 255])

# 设定红色阈值 红色有两种阈值，H[0:10] or H[156, 143]
lower_red_2 = np.array([156, 43, 43])
upper_red_2 = np.array([180, 255, 255])

# 设定黄色阈值
lower_yellow = np.array([11, 43, 43])
upper_yellow = np.array([40, 255, 255])


"""
一个基于颜色分割的交通灯识别程序
"""
class TrafficRec():
    def __init__(self):
        self.traffic = 0  # 识别结果

    def getRec(self):
        return self.traffic

    def colorDivision(self, frame, lower, upper):
        hsv = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
        mask = cv.inRange(hsv, lower, upper)
        # contours, hierarchy = cv.findContours(mask, cv.RETR_TREE, \
        #                                       cv.CHAIN_APPROX_SIMPLE)[-2:]
        return cv.bitwise_and(frame, frame, mask = mask)

    def maxArea(self, contours):
        area_max = 0
        for cnt in contours:
            Area = cv.contourArea(cnt)
            if Area > area_max:
                area_max = Area
        return area_max

    def recTraffic(self, frame):
        h, w = frame.shape[:2]
        roi = frame[int(h * .1):int(h * .6), int(.1 * w):int(.7 * w)]
        cv.imshow("roi", roi)
        gray = cv.cvtColor(roi, cv.COLOR_BGR2GRAY)
        ret, thredshold = cv.threshold(gray, 0, 255, cv.THRESH_OTSU + cv.THRESH_BINARY)
        contours, hierarchy = cv.findContours(thredshold, cv.RETR_TREE, \
                                              cv.CHAIN_APPROX_SIMPLE)[-2:]
        cnt_t = None
        area_max = 0
        for cnt in contours:
            Area = cv.contourArea(cnt)
            if Area > area_max:
                area_max = Area
                cnt_t = cnt

        if cnt_t is None and area_max < 100:
            return
        x, y, w, h = cv.boundingRect(cnt_t)
        _roi = roi[y:y + h, x:x + w]
        red = self.colorDivision(_roi, lower_red_2, upper_red_2)
        yellow = self.colorDivision(_roi, lower_yellow, upper_yellow)
        green = self.colorDivision(_roi, lower_green, upper_green)

        red_gray = cv.cvtColor(red, cv.COLOR_BGR2GRAY)
        yellow_gray = cv.cvtColor(yellow, cv.COLOR_BGR2GRAY)
        green_gray = cv.cvtColor(green, cv.COLOR_BGR2GRAY)

        _, red_thr = cv.threshold(red_gray, 40, 255, cv.THRESH_BINARY)
        _, yellow_thr = cv.threshold(yellow_gray, 40, 255, cv.THRESH_BINARY)
        _, green_thr = cv.threshold(green_gray, 40, 255, cv.THRESH_BINARY)

        # Rectangular Kernel
        rectKernel = cv.getStructuringElement(cv.MORPH_CROSS, (3, 3))

        red_thr = cv.morphologyEx(red_thr, cv.MORPH_OPEN, rectKernel)
        yellow_thr = cv.morphologyEx(yellow_thr, cv.MORPH_OPEN, rectKernel)
        green_thr = cv.morphologyEx(green_thr, cv.MORPH_OPEN, rectKernel)

        cv.imshow('red', red_thr)
        cv.imshow('yellow', yellow_thr)
        cv.imshow('green', green_thr)

        red_contours, _ = cv.findContours(red_thr, cv.RETR_TREE, \
                                          cv.CHAIN_APPROX_SIMPLE)[-2:]

        yellow_contours, _ = cv.findContours(yellow_thr, cv.RETR_TREE, \
                                             cv.CHAIN_APPROX_SIMPLE)[-2:]

        green_contours, _ = cv.findContours(green_thr, cv.RETR_TREE, \
                                            cv.CHAIN_APPROX_SIMPLE)[-2:]
        red_area = self.maxArea(red_contours)
        yellow_area = self.maxArea(yellow_contours)
        green_area = self.maxArea(green_contours)
        maxVal = max(red_area, yellow_area, green_area)
        # 0X01（红灯）
        # 0X02（黄灯）
        # 0X03（绿灯）
        rec_str = ""
        if maxVal == red_area:
            self.traffic = 0x01
            rec_str = "红灯"
        elif maxVal == yellow_area:
            self.traffic = 0x02
            rec_str = "黄灯"
        else:
            self.traffic = 0x03
            rec_str = "绿灯"

        frame = putText(frame, rec_str)
        return frame
